import numpy as np
import matplotlib.pyplot as plt

CONST_16QAM = np.array([
    -3 + 1j * 3,
    -3 + 1j * 1,
    -3 + 1j * -3,
    -3 + 1j * -1,
    -1 + 1j * 3,
    -1 + 1j * 1,
    -1 + 1j * -3,
    -1 + 1j * -1,
    3 + 1j * 3,
    3 + 1j * 1,
    3 + 1j * -3,
    3 + 1j * -1,
    1 + 1j * 3,
    1 + 1j * 1,
    1 + 1j * -3,
    1 + 1j * -1,
] / np.sqrt(10))

def sym_decide_np(symbols, constellations):
    """[for numpy]
    Make decisions of the input symbols. Pick the nearest ideal symbol
    according to the value of input symbols.

    Parameters
    ----------
    symbols : numpy array
        Input symbols array.
    constellations : numpy array
        Constellations of the ideal symbols.

    Returns
    -------
    numpy array
        Symbol decisions.
    """
    const = constellations
    const = const.reshape(*[1] * symbols.ndim, -1)
    symbols = symbols / \
        (np.sqrt(np.mean(symbols*symbols.conj(), axis=-1, keepdims=True)) + 1e-20)
    symbol_dec_indx = np.argmin(np.abs(symbols[..., np.newaxis] - const),
                                axis=-1).astype(np.uint8)

    symbol_dec = np.take(const, symbol_dec_indx)

    return symbol_dec

def snr_dd(
        sig,
        constellations=CONST_16QAM,
        dB_out=True,
        power_norm=True,
):
    sig_ideal = sym_decide_np(sig, constellations)
    if power_norm:
        sig = sig / np.sqrt(np.mean(abs(sig)**2, axis=-1, keepdims=True))
    ret = np.sum(np.abs(sig_ideal)**2, axis=-1) / np.sum(
            np.abs(sig - sig_ideal)**2, axis=-1)
    if dB_out == True:
        ret = 10 * np.log10(ret)
    return ret

def adaptive_MZM_bias_compensate(sig,
                                 constellations=CONST_16QAM,
                                 **opt_kwargs):
    """[summary]
        补偿IQ不匹配
    Parameters
    ----------
    sig : [type] np.ndarray
        [description] 经过线性均衡的输入信号，1倍采样，shape: [*, samples]
    constellations : [type], optional, np.ndarray
        [description] 星座图分布，星座点顺序无所谓

    Returns
    -------
    [type]
        [description] 补偿后的信号
    """

    # Bayesian optimization lib: https://github.com/fmfn/BayesianOptimization
    from bayes_opt import BayesianOptimization

    def cal_snr(sig, phi):
        sig_recover = recover_sig_by_phi(sig, phi)
        snr = snr_dd(sig_recover, constellations)
        snr = np.mean(snr)
        return snr

    def recover_sig_by_phi(sig, phi):
        # sig_IQ_mismatch = g1*sig + g2*conj(sig)
        # sig = 1/(conj(g2)/conj(g1) - g1/g2) * (conj(sig_IQ_mismatch/g1) - sig_IQ_mismatch/g2)
        g1 = (1 + np.exp(1j * phi)) / 2
        g2 = (1 - np.exp(1j * phi)) / 2

        if g2 != 0:
            sig_recover = (1 / (g2.conj() / g1.conj() - g1 / g2)) * (
                (sig / g1).conj() - sig / g2)
        else:
            return sig
        return sig_recover

    # 贝叶斯优化要最大化的目标函数
    obj_func = lambda phi: cal_snr(sig, phi)

    # 配置贝叶斯优化器
    phi_bound = opt_kwargs.get('phi_bound',
                               {'phi': (-np.pi / 2 * 0.5, np.pi / 2 * 0.5)})
    verbose = opt_kwargs.get('verbose', 2)
    optimizer = BayesianOptimization(
        f=obj_func,
        pbounds=phi_bound,
        verbose=verbose,
    )

    optimizer.maximize(init_points=5, n_iter=20, acq='ei', xi=0)
    phi_opt = optimizer.max['params']['phi']
    return recover_sig_by_phi(sig, phi_opt)


def decide_indx(symbols, constellations=CONST_16QAM, mod_order=4):
    ''' symbol.shape = [*, samples] '''
    constellation_points = constellations
    constellation_points = constellation_points.reshape(
        *[1] * symbols.ndim, -1)
    symbols = symbols / (np.sqrt(
        np.mean(symbols * symbols.conj(), axis=-1, keepdims=True)) + 1e-16)
    symbol_dec = np.argmin(np.abs(symbols[..., np.newaxis] -
                                  constellation_points),
                           axis=-1).astype(np.uint8)
    return symbol_dec

def map(prbs, constellations=CONST_16QAM, mod_order=4, order='msb'):

    base_vec = np.arange(mod_order)
    base_vec = 2**base_vec
    if order == 'msb':
        base_vec = np.flip(base_vec)
    elif order == 'lsb':
        pass
    else:
        raise Exception('Invalid bit order.')

    # prbs_len = prbs.shape[-1]
    # prbs_reshape = prbs.reshape()

    symbol_dec = np.sum(prbs.reshape(*prbs.shape[0:-1], -1, mod_order) *
                        base_vec.reshape(*[1] * prbs.ndim, -1),
                        axis=-1)
    ret = constellations[symbol_dec]
    return ret

def demap(symbols, constellations=CONST_16QAM, mod_order=4, order='msb'):
    ''' symbol.shape = [*, samples] '''
    ''' Only for 16QAM '''
    # constellation_points = constellations
    # constellation_points = constellation_points.reshape(*[1]*symbols.ndim, -1)
    # symbols=symbols/(np.sqrt(np.mean(symbols*symbols.conj(), axis=-1, keepdims=True))+1e-16)
    # symbol_dec = np.argmin(np.abs(
    #     symbols[..., np.newaxis]-constellation_points), axis=-1).astype(np.uint8)
    # ret = np.flip(np.unpackbits(
    #     symbol_dec[..., np.newaxis], count=-mod_order, axis=-1, bitorder='little'), axis=-1)
    symbol_dec = decide_indx(symbols, constellations, mod_order)
    bits = np.unpackbits(symbol_dec[..., np.newaxis],
                         axis=-1,
                         bitorder='little')

    if order == 'msb':
        ret = np.flip(bits[..., :mod_order], axis=-1)
    elif order == 'lsb':
        ret = bits[..., :mod_order]
    else:
        raise Exception('Invalid order, which is supposed to be msb or lsb.')
    ret = ret.reshape(*ret.shape[0:-2], ret.shape[-1] * ret.shape[-2])
    return ret

def test_case():
    prbs = np.random.randint(low=0, high=2, size=(1, 4096))
    phi = np.pi / 2*0.1

    g1 = (1 + np.exp(1j * phi)) / 2
    g2 = (1 - np.exp(1j * phi)) / 2

    sig_tx = map(prbs=prbs)

    plt.plot(sig_tx.real, sig_tx.imag, 'x')
    plt.title('Original signal')
    plt.show()
    plt.close()

    sig_iq_ib = sig_tx * g1 + sig_tx.conj() * g2
    plt.plot(sig_iq_ib.real, sig_iq_ib.imag, 'x')
    plt.title('IQ mismatch')
    plt.close()


    sig_recover = adaptive_MZM_bias_compensate(sig_iq_ib)

    plt.plot(sig_recover.real, sig_recover.imag, 'x' )
    plt.title('Recovered signal')
    plt.close()
